from model_train import roberta_chinese
import pandas as pd
import numpy as np

def pre_process():
  df =pd.read_csv('0208_0313_merge.csv')
  conditions = [
    df["content"] == 1,
    df["content"] != 1
  ]
  choices = [1, 0]
  df["label"] = np.select(conditions, choices, default=0)
  df.to_csv('0208_0313_merge_2.csv', index=False)

def tune_2nd(df, df_old, model_name, num_labels, model_save_path, content_column, label_column):
  roberta_chinese.tune_model_classify(df, df_old, model_name, num_labels, model_save_path, content_column, label_column)


if __name__ == "__main__":
  # pre_process()
  df =pd.read_csv('0208_0313_merge.csv')
  df_old = pd.read_csv('ChnSentiCorp_htl_all.csv')
  df_old = df_old.rename(columns={'review':'summary'})
  df_old = df_old.rename(columns={'label':'content'})
  model_name = 'model/roberta-tune-htl'
  num_labels = 3
  model_save_path = '/Users/zhangxiaotian/PycharmProjects/pytorch-bert/model/roberta-tune-htl-summary-3num'
  content_column = 'summary'
  label_column = 'content'
  tune_2nd(df, df_old, model_name, num_labels, model_save_path, content_column, label_column)
  # roberta_chinese.tune_model(df, model_name, num_labels, model_save_path, content_column, label_column)
